package cn.youfule.rtsp2ws;


import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.BinaryMessage;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.BinaryWebSocketHandler;

import java.util.Map;
import java.util.Set;

@Component
public class WsHandler extends BinaryWebSocketHandler {

    @Autowired
    private ConvertVideoPakcet convertVideoPakcet;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        System.out.println("新的加入"+session.getId());

        Map<String,Object> attr = session.getAttributes();
        Object sourceObj = attr.get(CommonUtil.PARAM_SOURCE);
        if (sourceObj != null) {
            String source = sourceObj.toString();

            //获取视频大小参数
            Object rsObj = attr.get(CommonUtil.PARAM_RS);
            String rs = rsObj != null ? rsObj.toString() : CommonUtil.PARAM_RS_DEFAULT;

            //获取sourceKey
            String sourceKey = CommonUtil.getSourceKey(source,rs);
            // 用户连接成功，放入在线用户缓存
            WsSessionManager.add(session.getId(),sourceKey, session);
            // 推送该IP对应的流
            convertVideoPakcet.startPush(source,rs);
        } else {
            throw new RuntimeException("视频连接失败!");
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        System.out.println("退出"+session.getId());

        this.stopSend(session);
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        System.out.println("传输错误"+session.getId());

        this.stopSend(session);
    }

    /**
     * 删除session缓存，并且停止推流
     * @param session
     */
    private void stopSend(WebSocketSession session){
        Map<String,Object> attr = session.getAttributes();
        Object sourceObj = attr.get(CommonUtil.PARAM_SOURCE);
        if (sourceObj != null) {
            String source = sourceObj.toString();

            //获取视频大小参数
            Object rsObj = attr.get(CommonUtil.PARAM_RS);
            String rs = rsObj != null ? rsObj.toString() : CommonUtil.PARAM_RS_DEFAULT;

            //获取sourceKey
            String sourceKey = CommonUtil.getSourceKey(source,rs);
            // 删除缓存的session
            WsSessionManager.removeAndClose(session.getId(),sourceKey);
            // 不存在客户端则停止推流
            Set<String> keys = WsSessionManager.getKeys(sourceKey);
            if(keys == null || keys.size() == 0){
                convertVideoPakcet.stopPush(source,rs);
            }
        } else {
            throw new RuntimeException("视频连接失败!");
        }
    }


    /**
     * 发送视频流到前端
     * @param sourceKey
     * @param data
     */
    public void sendVideo(String sourceKey,byte[] data) {
        BinaryMessage binaryMessage = new BinaryMessage(data);
        Set<WebSocketSession> clients = WsSessionManager.getSessions(sourceKey);
        for (WebSocketSession session: clients) {
            try {
                if (session.isOpen()) {
                    session.sendMessage(binaryMessage);
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }
}
